import tensorflow as tf

# 每次训练同一张图会生成两个不同版本，以便模型学习“相同语义”的不变性。
def random_augment(image, image_size=(640, 640)):
    # 随机裁剪 + resize
    image = tf.image.random_crop(image, size=(int(image_size[0] * 0.8), int(image_size[1] * 0.8), 3))
    image = tf.image.resize(image, image_size)

    # 随机翻转
    image = tf.image.random_flip_left_right(image)

    # 随机颜色扰动
    image = tf.image.random_brightness(image, max_delta=0.3)
    image = tf.image.random_contrast(image, lower=0.6, upper=1.4)
    image = tf.image.random_saturation(image, lower=0.6, upper=1.4)
    image = tf.image.random_hue(image, max_delta=0.1)

    # 保证像素值在 0~1
    image = tf.clip_by_value(image / 255.0, 0.0, 1.0)
    return image

# 构造 SimSiam 训练的核心机制：两个视图（view1/view2）
def simsiam_augment(image, image_size=(640, 640)):
    view1 = random_augment(image, image_size)
    view2 = random_augment(image, image_size)
    return view1, view2
